package com.wyg.demo.common.utils;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.csv.CSVFormat;
import org.apache.commons.csv.CSVPrinter;

import javax.servlet.http.HttpServletResponse;
import java.io.*;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.List;

/**
 * @author ygWen
 * @Description: 导出至csv文件
 * @date 2021/12/31
 */
@Slf4j
public class CsvUtil {
    /**
     * 写CSV并转换为字节流
     * @param headers 表头
     * @param cellList 表数据
     * @return
     */
    public static <T> byte[] writeCsvAfterToBytes(String[] headers, List<T[]> cellList) {
        byte[] bytes = new byte[0];
        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
        OutputStreamWriter outputStreamWriter = new OutputStreamWriter(byteArrayOutputStream, StandardCharsets.UTF_8);
        BufferedWriter bufferedWriter = new BufferedWriter(outputStreamWriter);
        CSVPrinter csvPrinter = null;
        try {
            //创建csvPrinter并设置表格头
            csvPrinter = new CSVPrinter(bufferedWriter, CSVFormat.DEFAULT.withHeader(headers));
            //写数据
            csvPrinter.printRecords(cellList);
            csvPrinter.flush();
            bytes = byteArrayOutputStream.toString(StandardCharsets.UTF_8.name()).getBytes();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                if (csvPrinter != null) {
                    csvPrinter.close();
                }
                bufferedWriter.close();
                outputStreamWriter.close();
                byteArrayOutputStream.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return bytes;
    }

    /**
     * 设置下载响应
     * @param fileName
     * @param bytes
     * @param response
     */
    public static void responseSetProperties(String fileName, byte[] bytes, HttpServletResponse response) {
        try {
            fileName = URLEncoder.encode(fileName, StandardCharsets.UTF_8.name());
            response.setContentType("application/csv");
            response.setCharacterEncoding(StandardCharsets.UTF_8.name());
            response.setHeader("Pragma", "public");
            response.setHeader("Cache-Control", "max-age=30");
            response.setHeader("Content-Disposition", "attachment; filename=" + fileName);
            OutputStream outputStream = response.getOutputStream();
            outputStream.write(bytes);
            outputStream.flush();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}
